import numpy as np
import matplotlib.pyplot as plt
import os

model_name = 'cnn_sim2024-11-23_16-00'
sample_num = 5
freq_num = 151

model_output_array = np.load(f'test/test_out_x{sample_num}_{model_name}.npy')
res_array = np.load(f'test/CST_out_x{sample_num}_{model_name}.npy')

freq_list = np.linspace(150, 300, freq_num)
s_para_name_list = ['tyx', 'txx', 'ryx', 'rxx']

try:
    os.makedirs(f'test/fig/temp/{model_name}')
except FileExistsError:
    pass

for sample_idx in range(sample_num):
    sample_fig_path = f'test/fig/temp/{model_name}/sample{sample_idx}'

    try:
        os.makedirs(sample_fig_path)
    except FileExistsError:
        pass

    for s_para_idx in range(4):
        plt.figure(figsize=(8, 4), dpi=200)
        plt.ylim(-1, 1)

        plt.plot(freq_list,
                 model_output_array[sample_idx, s_para_idx * freq_num * 2: (s_para_idx * 2 + 1) * freq_num],
                 '-', label='NN_real')
        plt.plot(freq_list,
                 model_output_array[sample_idx, s_para_idx * freq_num * 2 + freq_num: (s_para_idx * 2 + 2) * freq_num],
                 '-', label='NN_imag')
        plt.plot(freq_list,
                 res_array[sample_idx, s_para_idx * freq_num * 2: (s_para_idx * 2 + 1) * freq_num],
                 '--', label='CST_real')
        plt.plot(freq_list,
                 res_array[sample_idx, s_para_idx * freq_num * 2 + freq_num: (s_para_idx * 2 + 2) * freq_num],
                 '--', label='CST_imag')
        plt.legend()
        plt.xlabel('Freq(THz)')
        plt.ylabel(s_para_name_list[s_para_idx])
        plt.tight_layout()
        plt.savefig(f'{sample_fig_path}/{s_para_name_list[s_para_idx]}.png')
